from distilabel.llms import InferenceEndpointsLLM
from distilabel.pipeline import Pipeline
from distilabel.steps.tasks.improving_text_embeddings import (
    BitextRetrievalGenerator,
    EmbeddingTaskGenerator,
    GenerateLongTextMatchingData,
    GenerateShortTextMatchingData,
    GenerateTextClassificationData,
    GenerateTextRetrievalData,
    MonolingualTripletGenerator,
)

with Pipeline(name="improving-text-embeddings-with-llms") as pipeline:
    brainstorm_retrieval = EmbeddingTaskGenerator(
        category="text-retrieval",
        flatten_tasks=True,
        llm=InferenceEndpointsLLM(
            model_id="CohereForAI/c4ai-command-r-plus",
            tokenizer_id="CohereForAI/c4ai-command-r-plus",
        ),
        num_generations=1,
        group_generations=True,
        output_mappings={"model_name": "brainstorm_model"},
    )

    generate_retrieval = GenerateTextRetrievalData(
        add_raw_output=True,
        llm=InferenceEndpointsLLM(
            model_id="CohereForAI/c4ai-command-r-plus",
            tokenizer_id="CohereForAI/c4ai-command-r-plus",
        ),
        output_mappings={"model_name": "generation_model"},
    )

    brainstorm_retrieval >> generate_retrieval  # type: ignore

    brainstorm_classification = EmbeddingTaskGenerator(
        category="text-classification",
        flatten_tasks=True,
        llm=InferenceEndpointsLLM(
            model_id="CohereForAI/c4ai-command-r-plus",
            tokenizer_id="CohereForAI/c4ai-command-r-plus",
        ),
        num_generations=1,
        group_generations=True,
        output_mappings={"model_name": "brainstorm_model"},
    )

    generate_classification = GenerateTextClassificationData(
        add_raw_output=True,
        llm=InferenceEndpointsLLM(
            model_id="CohereForAI/c4ai-command-r-plus",
            tokenizer_id="CohereForAI/c4ai-command-r-plus",
        ),
        output_mappings={"model_name": "generation_model"},
    )

    brainstorm_classification >> generate_classification  # type: ignore

    brainstorm_matching_short = EmbeddingTaskGenerator(
        category="text-matching-short",
        flatten_tasks=True,
        llm=InferenceEndpointsLLM(
            model_id="CohereForAI/c4ai-command-r-plus",
            tokenizer_id="CohereForAI/c4ai-command-r-plus",
        ),
        num_generations=1,
        group_generations=True,
        output_mappings={"model_name": "brainstorm_model"},
    )

    generate_matching_short = GenerateShortTextMatchingData(
        add_raw_output=True,
        llm=InferenceEndpointsLLM(
            model_id="CohereForAI/c4ai-command-r-plus",
            tokenizer_id="CohereForAI/c4ai-command-r-plus",
        ),
        output_mappings={"model_name": "generation_model"},
    )

    brainstorm_matching_short >> generate_matching_short  # type: ignore

    brainstorm_matching_long = EmbeddingTaskGenerator(
        category="text-matching-long",
        flatten_tasks=True,
        llm=InferenceEndpointsLLM(
            model_id="CohereForAI/c4ai-command-r-plus",
            tokenizer_id="CohereForAI/c4ai-command-r-plus",
        ),
        num_generations=1,
        group_generations=True,
        output_mappings={"model_name": "brainstorm_model"},
    )

    generate_matching_long = GenerateLongTextMatchingData(
        add_raw_output=True,
        llm=InferenceEndpointsLLM(
            model_id="CohereForAI/c4ai-command-r-plus",
            tokenizer_id="CohereForAI/c4ai-command-r-plus",
        ),
        output_mappings={"model_name": "generation_model"},
    )

    brainstorm_matching_long >> generate_matching_long  # type: ignore

    bitext_retrieval_generator = BitextRetrievalGenerator(
        llm=InferenceEndpointsLLM(
            model_id="CohereForAI/c4ai-command-r-plus",
            tokenizer_id="CohereForAI/c4ai-command-r-plus",
        ),
        output_mappings={"model_name": "bitext_model"},
    )

    monolingual_triplet_generator = MonolingualTripletGenerator(
        llm=InferenceEndpointsLLM(
            model_id="CohereForAI/c4ai-command-r-plus",
            tokenizer_id="CohereForAI/c4ai-command-r-plus",
        ),
        output_mappings={"model_name": "monolingual_model"},
    )


if __name__ == "__main__":
    distiset = pipeline.run(
        parameters={
            step_name: {
                "llm": {
                    "generation_kwargs": {
                        "temperature": 0.7,
                        "max_new_tokens": 4096,
                        "stop_sequences": ["<EOS_TOKEN>", "<|END_OF_TURN_TOKEN|>"],
                    }
                }
            }
            for step_name in pipeline.dag
        },
    )
    if distiset is not None:
        distiset.push_to_hub(
            "distilabel-internal-testing/alvarobartt-improving-text-embeddings-with-llms-full",
        )